import datetime
import os.path
import typing as tp
import json
import traceback

from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.date import DateTrigger
from apscheduler.triggers.cron import CronTrigger
from apscheduler.triggers.interval import IntervalTrigger
from loguru import logger

from .agent import Agent
from .jobs.single_node_single_mod_job import SingleNSingleMJob
from .jobs.units import JobType
from .jobs import JobVarType
from .db.models.job import Job
from .common import Result, ResultItem
from .runners.snsm_runner import SNSMRunner
from .task import Task


class Scheduler:
    def __init__(self, agent: Agent):
        self.agent = agent
        self._aps_sche = AsyncIOScheduler()

    @property
    def scheduler(self):
        return self._aps_sche

    def start(self):
        self._aps_sche.start()

    def shutdown(self):
        self._aps_sche.shutdown()


class JobScheduler(Scheduler):
    def __init__(self, agent: Agent):
        super().__init__(agent)
        self.id2job = dict()

    async def create_snsm_job(self, job: SingleNSingleMJob) -> tp.Tuple[bool, Result]:
        start_date = job.sche.start_date or datetime.datetime.now()
        end_date = job.sche.end_date

        obj = Job(
            name=job.name,
            creator=job.creator,
            send_list=json.dumps(job.send_list),
            cc_list=json.dumps(job.cc_list),

            aps_type=job.sche.type,
            aps_cfg=job.sche.model_dump_json(),

            step=job.first_step.model_dump_json(),

            is_unique=job.is_unique,
            type=JobType.snsm,
            tag=job.tag
        )
        if job.sche.type == "date":
            trigger = DateTrigger(**job.sche.date.model_dump())
            if job.sche.date.run_date:
                start_date = job.sche.date.run_date
            end_date = start_date
        elif job.sche.type == "interval":
            trigger = IntervalTrigger(start_date=start_date, end_date=end_date, *job.sche.interval.model_dump())
        else:
            trigger = CronTrigger(start_date=start_date, end_date=end_date, **job.sche.corn.model_dump())
        obj.start_date = start_date
        obj.end_date = end_date

        try:
            await self.agent.db.add(obj)
            aps_job = self._aps_sche.add_job(
                id=str(obj.id),
                func=self.run_task,
                trigger=trigger,
                args=(str(obj.id), obj.type, job),
                kwargs={},
                next_run_time=start_date
            )
            self.id2job[str(obj.id)] = aps_job
            logger.info(f"[{obj.id}] {job} 调度配置：{trigger}，首次触发时间: {aps_job.next_run_time}")
        except (BaseException, Exception):
            import traceback
            logger.error(traceback.format_exc())
            return False, Result.FRAMEWORK_ERROR
        else:
            return True, Result.SUCCESS

    async def run_task(self, id_: str, type_: JobType, job: JobVarType):
        if type_ == JobType.snsm:
            task = await self.agent.task_scheduler.create_task(id_, job)
            runner = SNSMRunner(self.agent, task)
            task.runner = runner
            await runner.run()
            task.runner = None
            await self.agent.task_scheduler.remove_task(id_)

    async def remove_job(self, id_: str) -> bool:
        if id_ in self.id2job:
            self.id2job.pop(id_)
            self._aps_sche.remove_job(id_)
            return True
        else:
            return False


class TaskScheduler(Scheduler):
    def __init__(self, agent: Agent):
        super().__init__(agent)
        self.id2task = dict()

    async def create_task(self, id_: str, job: JobVarType) -> Task:
        task = Task(
            job=job,
            aps_job_id=id_,
            log_root_path=os.path.join(self.agent.cfg.task_path, id_)
        )
        self.id2task[id_] = task
        return task

    async def remove_task(self, id_: str) -> bool:
        if id_ in self.id2task:
            self.id2task.pop(id_)
            return True
        else:
            return False

    async def interrupt(self, job_id: int) -> ResultItem:
        job_id = str(job_id)
        if job_id in self.id2task:
            task = self.id2task.pop(job_id)
            logger.info(f"[{task.aps_job_id}] 收到任务中断请求")
            await task.runner.interrupt()
            return ResultItem(msg="任务中断请求已处理", code=0)
        else:
            return ResultItem(msg=f"未找到任务", code=1)
